import logging

import torch
from saicinpainting.training.trainers.base import BaseInpaintingTrainingModule

LOGGER = logging.getLogger(__name__)


class DefaultInpaintingTrainingModule(BaseInpaintingTrainingModule):
    def __init__(self, *args, concat_mask=True, rescale_scheduler_kwargs=None, image_to_discriminator='predicted_image',
                 add_noise_kwargs=None, noise_fill_hole=False, const_area_crop_kwargs=None,
                 distance_weighter_kwargs=None, distance_weighted_mask_for_discr=False,
                 fake_fakes_proba=0, fake_fakes_generator_kwargs=None,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.concat_mask = concat_mask

    def forward(self, batch):

        img = batch['image']
        mask = batch['mask']
        rel_pos = batch['rel_pos']
        direct = batch['direct']

        masked_img = img * (1 - mask)
        
        if self.concat_mask:
            masked_img = torch.cat([masked_img, mask], dim=1)

        batch['predicted_image'] = self.generator(masked_img, rel_pos, direct)
        batch['inpainted'] = mask * batch['predicted_image'] + (1 - mask) * batch['image']

        return batch